#include "lib/print.h"
#include "trap/trap.h"
#include "proc/cpu.h"
#include "mem/vmem.h"
#include "memlayout.h"
#include "riscv.h"

// in trampoline.S
extern char trampoline[];      // 内核和用户切换的代码
extern char user_vector[];     // 用户触发trap进入内核
extern char user_return[];     // trap处理完毕返回用户

// in trap.S
extern char kernel_vector[];   // 内核态trap处理流程

// in trap_kernel.c
extern char* interrupt_info[16]; // 中断错误信息
extern char* exception_info[16]; // 异常错误信息

// 在user_vector()里面调用
// 用户态trap处理的核心逻辑
void trap_user_handler()
{
    uint64 sepc = r_sepc();          // 记录了发生异常时的pc值
    uint64 sstatus = r_sstatus();    // 与特权模式和中断相关的状态信息
    uint64 scause = r_scause();      // 引发trap的原因
    uint64 stval = r_stval();        // 发生trap时保存的附加信息(不同trap不一样)
    proc_t* p = myproc();

    // 确认trap来自U-mode
    assert((sstatus & SSTATUS_SPP) == 0, "trap_user_handler: not from u-mode");

    // 切换陷阱向量到内核态的处理程序
    w_stvec((uint64)kernel_vector);

    // 保存当前陷阱的 EPC 到进程的 trapframe
    p->tf->epc = sepc;

    // 低四位的值作为trap类型的编码信息
    int trap_id = scause & 0xf; 

     // 判断是否为中断
    if (scause & (1L << 63)) { // 中断类型的高位标志
        switch (trap_id) {
            case 1: // S-mode software interrupt
                timer_interrupt_handler();
                break;
            case 9: // S-mode external interrupt
                external_interrupt_handler();
                break;
            default:
                printf("Interrupt: %s\n", interrupt_info[trap_id]);
                printf("Trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                panic("Unexpected interrupt in trap_user_handler");
                break;
        }
    } else { // 异常处理
        switch (trap_id) {
            case 8: // 系统调用（环境调用）
                p->tf->epc += 4; // 系统调用指令长度为 4 字节，跳过当前指令
                intr_on();       // 开启中断（允许外部中断继续处理）
                printf("get a syscall from proc %d\n", p->pid);
                // 系统调用处理可以在这里扩展
                break;
            default:
                if (trap_id < 16) {
                    printf("Exception: %s\n", exception_info[trap_id]);
                } else {
                    printf("Unexpected exception: trap_id=%d\n", trap_id);
                }
                printf("Trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                panic("Unexpected exception in trap_user_handler");
                break;
        }
    }
    // 调用user_return()
    // 内核态返回用户态
    trap_user_return();
}

// 调用user_return()
// 内核态返回用户态
void trap_user_return()
{
    proc_t *p = myproc(); 

    // 这里关闭中断是为了切换陷阱处理的目标
    intr_off();

    // 设置陷阱向量只想uservec
    // 计算地址：TRAMPOLINE 的基址加上 user_vector 在 trampoline 中的偏移量
    w_stvec(TRAMPOLINE + (user_vector - trampoline));

    // 逐项处理内核相关关键字，用于在用户态发生陷阱时返回到内核态
    p->tf->kernel_satp = r_satp();         // 保存当前的内核页表基址（satp 寄存器）
    p->tf->kernel_sp = p->kstack + PGSIZE; // 设置内核栈顶指针（指向分配的内核栈顶部）
    p->tf->kernel_trap = (uint64)trap_user_handler; // 设置用户态陷阱处理程序的入口地址
    p->tf->kernel_hartid = r_tp(); // 保存 hartid（当前 CPU 核心的 ID）

    // 修改 sstatus 寄存器：
    // 清除 SPP 位，确保从用户态返回（SPP = 0 表示用户态）。
    // 设置 SPIE 位，确保在返回到用户态后打开中断。
    unsigned long x = r_sstatus();
    x &= ~SSTATUS_SPP;  // 清除 Supervisor Previous Privilege 位
    x |= SSTATUS_SPIE;  // 设置 Supervisor Previous Interrupt Enable 位
    w_sstatus(x);       // 写回修改后的 sstatus

    // 设置 sepc 寄存器为用户程序的入口地址（epc）
    w_sepc(p->tf->epc);

    // 准备切换到用户态，设置用户页表基址到 satp
    uint64 satp = MAKE_SATP(p->pgtbl);

    // 计算用户态返回函数 user_return 的地址：
    // TRAMPOLINE 的基址加上 user_return 在 trampoline 中的偏移量
    uint64 fn = TRAMPOLINE + (user_return - trampoline);

    // 通过 trampoline 的跳板代码切换到用户态
    // 参数：
    //  - TRAPFRAME：指向用户态的 trapframe（保存了用户态的寄存器状态）
    //  - satp：用户页表基址，用于切换地址空间
    ((void (*)(uint64, uint64))fn)(TRAPFRAME, satp);
}